# websocket的socket实现，来源csdn，已魔改
# https://blog.csdn.net/difu0201/article/details/101582682
import sys
try:
    import ubinascii as binascii
except:
    import binascii
try:
    import uhashlib as hashlib
except:
    import hashlib
def server_handshake(sock):
    clr = sock.makefile("rwb", 0)
    l = clr.readline()
    webkey = None
    while 1:
        l = clr.readline()
        if not l:
            raise OSError("EOF in headers")
        if l == b"\r\n":
            break
        h, v = [x.strip() for x in l.split(b":", 1)]
        if DEBUG:
            print((h, v))
        if h == b'Sec-WebSocket-Key':
            webkey = v
    if not webkey:
        raise OSError("Not a websocket request")
    respkey = webkey + b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    respkey = hashlib.sha1(respkey).digest()
    respkey = binascii.b2a_base64(respkey)[:-1]
    resp = b"""\
HTTP/1.1 101 Switching Protocols\r
Upgrade: websocket\r
Connection: Upgrade\r
Sec-WebSocket-Accept: %s\r
\r
""" % respkey
    sock.send(resp)
def client_handshake(sock):
    cl = sock.makefile("rwb", 0)
    cl.write(b"""\
GET / HTTP/1.1\r
Host: echo.websocket.org\r
Connection: Upgrade\r
Upgrade: websocket\r
Sec-WebSocket-Key: foo\r
\r
""")
    l = cl.readline()
    while 1:
        l = cl.readline()
        if l == b"\r\n":
            break
        
import socket
try:
    import network
except:
    pass
import sys
import os
try:
    import ustruct as struct
except:
    import struct
    
DEBUG = False

class websockets:

    def __init__(self, s):
        self.s = s
        
    def write(self, data):
        l = len(data)
        if l < 126:
            hdr = struct.pack(">BB", 0x82, l)
        else:
            hdr = struct.pack(">BBH", 0x82, 126, l)
        self.s.send(hdr)
        self.s.send(data)

    def recvexactly(self, sz):
        res = b""
        while sz:
            data = self.s.recv(sz)
            if not data:
                break
            res += data
            sz -= len(data)
        return res

    def read(self):
        while True:
            hdr = self.recvexactly(2)
            assert len(hdr) == 2
            firstbyte, secondbyte = struct.unpack(">BB", hdr)

            mskenable =  True if secondbyte & 0x80 else False
            length = secondbyte & 0x7f
            if length == 126:
                hdr = self.recvexactly(2)
                assert len(hdr) == 2
                (length,) = struct.unpack(">H", hdr)
            if length == 127:
                hdr = self.recvexactly(8)
                assert len(hdr) == 8
                (length,) = struct.unpack(">Q", hdr)
            opcode =  firstbyte & 0x0f
            if opcode == 8:
                self.s.close()
                return ''
            fin = True if firstbyte&0x80 else False
            if mskenable:
                hdr = self.recvexactly(4)
                assert len(hdr) == 4
                (msk1,msk2,msk3,msk4) = struct.unpack(">BBBB", hdr)
                msk = [msk1,msk2,msk3,msk4]
            data = []
            while length:
                skip = self.s.recv(length)
                length -= len(skip)
                data.extend(skip)
            newdata = []
            for i,item in enumerate(data):
                j = i % 4
                newdata.append(chr(data[i] ^ msk[j]))
            res = ''.join(newdata)
            return res
class WebSocket:
    connect_max=1
    on_message=lambda s,a:a
    on_connect=lambda s,a:a
    on_close  =lambda s,a:a
    sendargs  =0
    def __init__(self,connect_max=1):
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.connect_max=connect_max
        
    def listen(self,ip_port,while_listen=False):
        self.sock.bind(ip_port)
        self.sock.listen(self.connect_max)
        while True:
            conn, address = self.sock.accept()
            server_handshake(conn)
            self.ws = websockets(conn)
            self.sendargs=address
            self.on_connect(address)
            while True:
                text = self.ws.read()
                if text =='':
                    break
                self.sendargs=text
                self.on_message(text)
            self.sendargs=address
            self.on_close(address)
            if not while_listen:
                break
    def sendMessage(self,data):
        self.ws.write(data)
    def onMessage(self,func):
        self.on_message=func
        def set_on_messgae():
            func(self.sendargs)
        return set_on_messgae
    def onConnect(self,func):
        self.on_connect=func
        def set_on_connect():
            func(self.sendargs)
        return set_on_connect
    def onClose(self,func):
        self.on_close=func
        def set_on_close():
            func(self.sendargs)
        return set_on_close




